from math import pi, log
import numpy as np
import mindspore
from mindspore import nn, ops, Tensor, Parameter
from mindspore.common.initializer import initializer
from typing import Union, Optional

# helper functions

def exists(val):
    return val is not None

def default(val, d):
    return val if exists(val) else d

def broadcat(tensors, dim=-1):
    broadcasted_tensors = ops.broadcast_tensors(tensors)
    return ops.concat(broadcasted_tensors, axis=dim)

def slice_at_dim(t, dim_slice: slice, dim: int):
    dim += t.ndim if dim < 0 else 0
    indices = [slice(None)] * t.ndim
    indices[dim] = dim_slice
    return t[tuple(indices)]

# rotary embedding helper functions

def rotate_half(x):
    x = ops.reshape(x, (*x.shape[:-1], -1, 2))
    x1, x2 = ops.split(x, axis=-1, output_num=2)
    x = ops.stack((-x2, x1), axis=-1)
    return ops.reshape(x, (*x.shape[:-2], -1))

def apply_rotary_emb(freqs, t, start_index=0, scale=1., seq_dim=-2, freqs_seq_dim=None):
    dtype = t.dtype

    if not exists(freqs_seq_dim):
        if freqs.ndim == 2 or t.ndim == 3:
            freqs_seq_dim = 0

    if t.ndim == 3 or exists(freqs_seq_dim):
        seq_len = t.shape[seq_dim]
        freqs = slice_at_dim(freqs, slice(-seq_len, None), dim=freqs_seq_dim)

    rot_dim = freqs.shape[-1]
    end_index = start_index + rot_dim

    assert rot_dim <= t.shape[-1], f'feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}'

    t_left = t[..., :start_index]
    t_middle = t[..., start_index:end_index]
    t_right = t[..., end_index:]

    t_transformed = (t_middle * ops.cos(freqs) * scale) + (rotate_half(t_middle) * ops.sin(freqs) * scale)

    out = ops.concat((t_left, t_transformed, t_right), axis=-1)

    return out.astype(dtype)

# learned rotation helpers

def apply_learned_rotations(rotations, t, start_index=0, freq_ranges=None):
    if exists(freq_ranges):
        rotations = ops.einsum('...,...f->...f', rotations, freq_ranges)
        rotations = ops.reshape(rotations, (*rotations.shape[:-1], -1))

    rotations = ops.tile(rotations, (1, 2))
    return apply_rotary_emb(rotations, t, start_index=start_index)

# classes

class RotaryEmbedding(nn.Cell):
    def __init__(
        self,
        dim: int,
        custom_freqs: Optional[Tensor] = None,
        freqs_for: str = 'lang',
        theta: float = 10000,
        max_freq: int = 10,
        num_freqs: int = 1,
        learned_freq: bool = False,
        use_xpos: bool = False,
        xpos_scale_base: int = 512,
        interpolate_factor: float = 1.,
        theta_rescale_factor: float = 1.,
        seq_before_head_dim: bool = False,
        cache_if_possible: bool = True,
        cache_max_seq_len: int = 8192
    ):
        super().__init__()

        theta *= theta_rescale_factor ** (dim / (dim - 2))

        self.freqs_for = freqs_for

        if exists(custom_freqs):
            freqs = custom_freqs
        elif freqs_for == 'lang':
            freqs = 1. / (theta ** (ops.arange(0, dim, 2)[:(dim // 2)].astype(mindspore.float32) / dim))
        elif freqs_for == 'pixel':
            freqs = ops.linspace(1., max_freq / 2, dim // 2) * pi
        elif freqs_for == 'constant':
            freqs = ops.ones((num_freqs,), mindspore.float32)

        self.cache_if_possible = cache_if_possible
        self.cache_max_seq_len = cache_max_seq_len

        self.cached_freqs = Parameter(initializer('zeros', [cache_max_seq_len, dim]), requires_grad=False)
        self.cached_freqs_seq_len = 0

        self.freqs = Parameter(freqs, requires_grad=learned_freq)

        self.learned_freq = learned_freq

        self.dummy = Parameter(Tensor(0), requires_grad=False)

        self.seq_before_head_dim = seq_before_head_dim
        self.default_seq_dim = -3 if seq_before_head_dim else -2

        assert interpolate_factor >= 1.
        self.interpolate_factor = interpolate_factor

        self.use_xpos = use_xpos

        if not use_xpos:
            return

        scale = (ops.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim)
        self.scale_base = xpos_scale_base

        self.scale = Parameter(scale, requires_grad=False)
        self.cached_scales = Parameter(initializer('zeros', [cache_max_seq_len, dim]), requires_grad=False)
        self.cached_scales_seq_len = 0

        self.apply_rotary_emb = staticmethod(apply_rotary_emb)

    @property
    def device(self):
        return self.dummy.device

    def get_seq_pos(self, seq_len, device, dtype, offset=0):
        return (ops.arange(seq_len, dtype=dtype) + offset) / self.interpolate_factor

    def rotate_queries_or_keys(self, t, seq_dim=None, offset=0, scale=None):
        seq_dim = default(seq_dim, self.default_seq_dim)

        assert not self.use_xpos or exists(scale), 'you must use `.rotate_queries_and_keys` method instead and pass in both queries and keys, for length extrapolatable rotary embeddings'

        dtype, seq_len = t.dtype, t.shape[seq_dim]

        seq = self.get_seq_pos(seq_len, device=device, dtype=dtype, offset=offset)

        freqs = self(seq, seq_len=seq_len, offset=offset)

        if seq_dim == -3:
            freqs = ops.reshape(freqs, (freqs.shape[0], 1, freqs.shape[1]))

        return apply_rotary_emb(freqs, t, scale=default(scale, 1.), seq_dim=seq_dim)

    def rotate_queries_with_cached_keys(self, q, k, seq_dim=None, offset=0):
        dtype, seq_dim = q.dtype, default(seq_dim, self.default_seq_dim)

        q_len, k_len = q.shape[seq_dim], k.shape[seq_dim]
        assert q_len <= k_len

        q_scale = k_scale = 1.

        if self.use_xpos:
            seq = self.get_seq_pos(k_len, dtype=dtype, device=device)
            q_scale = self.get_scale(seq[-q_len:]).astype(dtype)
            k_scale = self.get_scale(seq).astype(dtype)

        rotated_q = self.rotate_queries_or_keys(q, seq_dim=seq_dim, scale=q_scale, offset=k_len - q_len + offset)
        rotated_k = self.rotate_queries_or_keys(k, seq_dim=seq_dim, scale=k_scale ** -1)

        return rotated_q.astype(q.dtype), rotated_k.astype(k.dtype)

    def rotate_queries_and_keys(self, q, k, seq_dim=None):
        seq_dim = default(seq_dim, self.default_seq_dim)

        assert self.use_xpos
        device, dtype, seq_len = q.device, q.dtype, q.shape[seq_dim]

        seq = self.get_seq_pos(seq_len, dtype=dtype, device=device)

        freqs = self.forward(seq, seq_len=seq_len)
        scale = self.get_scale(seq, seq_len=seq_len).astype(dtype)

        if seq_dim == -3:
            freqs = freqs.reshape((freqs.shape[0], 1, freqs.shape[1]))
            scale = scale.reshape((scale.shape[0], 1, scale.shape[1]))

        rotated_q = apply_rotary_emb(freqs, q, scale=scale, seq_dim=seq_dim)
        rotated_k = apply_rotary_emb(freqs, k, scale=scale ** -1, seq_dim=seq_dim)

        return rotated_q.astype(q.dtype), rotated_k.astype(k.dtype)

    def get_scale(self, t: Tensor, seq_len: Optional[int] = None, offset=0):
        assert self.use_xpos

        should_cache = (
            self.cache_if_possible and
            exists(seq_len) and
            (offset + seq_len) <= self.cache_max_seq_len
        )

        if should_cache and exists(self.cached_scales) and (seq_len + offset) <= self.cached_scales_seq_len:
            return self.cached_scales[offset:(offset + seq_len)]

        scale = 1.
        if self.use_xpos:
            power = (t - len(t) // 2) / self.scale_base
            scale = self.scale ** power.reshape((-1, 1))
            scale = np.tile(scale, (1, 2))

        if should_cache and offset == 0:
            self.cached_scales[:seq_len] = scale
            self.cached_scales_seq_len = seq_len

        return scale

    def get_axial_freqs(self, *dims):
        all_freqs = []

        for ind, dim in enumerate(dims):
            if self.freqs_for == 'pixel':
                pos = np.linspace(-1, 1, num=dim)
            else:
                pos = np.arange(dim)

            freqs = self.forward(pos, seq_len=dim)

            all_axis = [None] * len(dims)
            all_axis[ind] = slice(None)

            new_axis_slice = (Ellipsis, *all_axis, slice(None))
            all_freqs.append(freqs[new_axis_slice])

        all_freqs = np.broadcast_arrays(*all_freqs)
        return np.concatenate(all_freqs, axis=-1)

    def forward(self, t: Tensor, seq_len: Optional[int] = None, offset=0):
        should_cache = (
            self.cache_if_possible and
            not self.learned_freq and
            exists(seq_len) and
            self.freqs_for != 'pixel' and
            (offset + seq_len) <= self.cache_max_seq_len
        )

        if should_cache and exists(self.cached_freqs) and (offset + seq_len) <= self.cached_freqs_seq_len:
            return self.cached_freqs[offset:(offset + seq_len)]

        freqs = self.freqs

        freqs = np.einsum('..., f -> ... f', t.astype(freqs.dtype), freqs)
        freqs = np.tile(freqs, (1, 2))

        if should_cache and offset == 0:
            self.cached_freqs[:seq_len] = freqs
            self.cached_freqs_seq_len = seq_len

        return freqs
